#!/usr/bin/python
# -*- coding: UTF-8 -*-
from datetime import datetime

from PyQt5 import QtCore, QtGui, QtWidgets
from ui.video_window import Ui_MainWindow
from ultralytics import YOLO
import cv2
import os
from detection import Worker
from ui.classDiaLogController import view_class_window
import uuid


class view_window(QtWidgets.QMainWindow):
    """打开窗口"""

    def __init__(self):
        QtWidgets.QMainWindow.__init__(self)
        self.ui = Ui_MainWindow()
        # 读取配置文件
        self.ui.setupUi(self)
        self.model_list = None
        self.class_dialog = None
        self.ui.video_btn.clicked.connect(self.__get_video_file)
        self.ui.model_btn.clicked.connect(self.__get_model_file)
        self.ui.handle_type_inp.activated.connect(self.__handle_type_change)
        self.ui.begin.clicked.connect(self.__begin)
        self.ui.video_track.clicked.connect(self.__video_track)
        self.ui.stop.clicked.connect(self.__stop)
        self.ui.start_inp.editingFinished.connect(self.__text_changed_finshed)
        self.ui.class_btn.clicked.connect(self.__showClassDialog)
        self.frame_count = 0
        self.fps = 0
        self.duration = "0"
        self.worker = Worker()
        self.view_class = ViewClass()
        self.worker.image_loaded.connect(self.showImg)
        self.worker.cur_info_loaded.connect(self.showCurImgInfo)
        self.worker.image_info_loaded.connect(self.showImgInfo)
        self.view_class.show_cur_loaded.connect(self.showImg)

    def __stop(self):
        self.worker.stop()

    def __showClassDialog(self):
        if not self.model_list:
            self.__showMsg("请先选择模型！")
            return
        self.class_dialog = view_class_window()
        self.class_dialog.ok_signal.connect(self.__reload_model_list)
        self.class_dialog.loadData(self.model_list)

    def __reload_model_list(self, sub_model_list):
        self.model_list = sub_model_list
        self.__load_scroll_layout()

    def __showMsg(self, msg):
        QtWidgets.QMessageBox.critical(self, "错误", msg)

    def __text_changed_finshed(self):
        start_time = self.ui.start_inp.text()
        cur_frame = int(float(start_time) * self.fps)
        if cur_frame > self.frame_count:
            cur_frame = self.frame_count
        self.view_class.load_img(self, self.ui.video_inp.text(), cur_frame)

    def __start(self, check_type=1):
        if "" == self.ui.video_inp.text():
            self.__showMsg("请选择视频！")
        param = {
            "video_path": self.ui.video_inp.text(),
            "handle_type": self.ui.handle_type_inp.currentData(),
            "interval_num": self.ui.interval_num.text(),
            "start_time": self.ui.start_inp.text(),
            "end_time": self.ui.end_inp.text(),
            "model_list": self.model_list,
            "pause_time": self.ui.pause_inp.text(),
            "check_type": check_type
        }
        print("__start开始:", datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"))
        self.worker.load_data(self, param)

    def __begin(self):
        self.__start(1)

    def __video_track(self):
        if len(self.model_list) > 1:
            self.__showMsg("目标追踪模式只能选择一个模型！")
            return
        self.__start(2)

    def __handle_type_change(self, idx):
        _translate = QtCore.QCoreApplication.translate
        handle_type = self.ui.handle_type_inp.currentData()
        self.ui.start_Label.setText(_translate("MainWindow", "开始帧"))
        self.ui.end_Label.setText(_translate("MainWindow", "结束帧"))
        self.ui.end_inp.setText(_translate("MainWindow", str(self.frame_count)))
        if handle_type == "1":
            self.ui.interval_label.setText(_translate("MainWindow", "间隔时间"))
            self.ui.interval_unit_Label.setText(_translate("MainWindow", "(秒)"))
            self.ui.start_Label.setText(_translate("MainWindow", "开始时间"))
            self.ui.end_Label.setText(_translate("MainWindow", "结束时间"))
            self.ui.end_inp.setText(_translate("MainWindow", self.duration))
        elif handle_type == "2":
            self.ui.interval_label.setText(_translate("MainWindow", "提取间隔"))
            self.ui.interval_unit_Label.setText(_translate("MainWindow", "(帧)"))
        elif handle_type == "3":
            self.ui.interval_label.setText(_translate("MainWindow", "提取总数"))
            self.ui.interval_unit_Label.setText(_translate("MainWindow", "(张)"))

    @QtCore.pyqtSlot(QtGui.QPixmap)
    def showImg(self, pixmap):
        self.ui.img_label.setPixmap(pixmap)

    @QtCore.pyqtSlot(str)
    def showCurImgInfo(self, msg):
        self.ui.cur_img_label.setText(msg)

    @QtCore.pyqtSlot(str)
    def showImgInfo(self, msg):
        self.ui.img_info_label.setText("  识别到:{}".format(msg))

    def cvMatToQImage(self, cvMat):
        if len(cvMat.shape) == 2:
            # 灰度图是单通道，所以需要用Format_Indexed8
            rows, columns = cvMat.shape
            rows, columns = self.__re_size(rows, columns)
            cvMat = cv2.resize(cvMat, (rows, rows))
            bytesPerLine = columns
            return QtGui.QImage(cvMat.data, columns, rows, bytesPerLine, QtGui.QImage.Format_Indexed8)
        else:
            rows, columns, channels = cvMat.shape
            rows, columns = self.__re_size(rows, columns)
            cvMat = cv2.resize(cvMat, (columns, rows))
            bytesPerLine = channels * columns
            return QtGui.QImage(cvMat.data, columns, rows, bytesPerLine, QtGui.QImage.Format_RGB888).rgbSwapped()

    def __re_size(self, height, width):
        raw_height = 540
        raw_width = 960
        raw_ratio = raw_height / raw_width
        ratio = height / width
        # 如果ratio 比 raw_ratio 大,则以width为基础，否则以height
        if ratio > raw_ratio:
            new_ratio = raw_height / height
            new_width = width * new_ratio
            return int(raw_height), int(new_width)
        else:
            new_ratio = raw_width / width
            new_height = height * new_ratio
            return int(new_height), int(raw_width)

    def __get_video_file(self):
        get_filenames_path, ok = QtWidgets.QFileDialog.getOpenFileName(self, "选取视频文件", '',
                                                                       filter="视频文件(*.asf *.avi *.m4v *.mkv *.mov *.mp4 *.mpeg *.mpg *.ts *.wmv)")
        if ok:
            video_path = str(get_filenames_path)
            self.ui.video_inp.setText(video_path)
            video = cv2.VideoCapture(video_path)
            frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
            fps = video.get(cv2.CAP_PROP_FPS)
            self.duration = str(frame_count // fps)
            self.frame_count = frame_count
            self.fps = fps
            self.ui.end_inp.setText(self.duration)
            self.ui.video_info_label.setText(
                "  视频信息:\n\t文件地址:{}\n\t总帧数:{}\n\tfps:{}\n\t总时长:{}".format(video_path, str(frame_count),
                                                                                        str(int(fps)), self.duration))
            self.view_class.load_img(self, video_path, 0)

    def __load_scroll_layout(self):
        self.__cleanScrollLayout()
        total_num = 0
        for model_info in self.model_list:
            # model
            self.__add_scroll_label(total_num, model_info["uuid"], model_info["file_name"] + ":", -1)
            total_num = total_num + 1
            for class_info in model_info["class_list"]:
                self.__add_scroll_label(total_num, class_info["uuid"], class_info["class_name"], class_info["used"])
                total_num = total_num + 1

    def __get_model_file(self):
        get_filenames_path, ok = QtWidgets.QFileDialog.getOpenFileNames(self, "选取模型文件", '',
                                                                        filter="模型文件(*.pt)")
        if ok:
            print("加载模型开始:", datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"))
            self.model_list = []
            model_name_list = []
            for index, model_path in enumerate(get_filenames_path):
                model_file_name = os.path.basename(model_path)
                self.model_list.append(self.__get_class_by_model(model_path, model_file_name))
                model_name_list.append(model_file_name)
            self.ui.model_inp.setText(str(";".join(model_name_list)))
            self.__load_scroll_layout()
            print("加载模型完成:", datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"))

    def __add_scroll_label(self, total_num, objectName, labelText, type):
        label = QtWidgets.QLabel()
        label.setGeometry(QtCore.QRect(10, (total_num * 30) + 10, 80, 16))
        label.setObjectName(objectName)
        label.setText(QtCore.QCoreApplication.translate("MainWindow", ("  " if type >= 0 else "") + labelText))
        if type == -1:
            font = label.font()  # 获取当前字体
            font.setBold(True)  # 设置字体为粗体
            label.setFont(font)  # 应用新的字体设置
        elif type == 1:
            label.setStyleSheet("QLabel { color : green; }")
        elif type == 2:
            label.setStyleSheet("QLabel { color : blue; }")
        elif type == 3:
            label.setStyleSheet("QLabel { color : red; }")
        self.ui.scrollLayout.addWidget(label)

    def __get_class_by_model(self, model_path, model_file_name):
        # 加载模型分类
        model = YOLO(model_path)
        class_name_list = model.names
        class_list = []
        for idx in class_name_list:
            class_list.append({"class_name": class_name_list[idx], "uuid": str(uuid.uuid4()), "used": 0})
        model_info = {"file_name": model_file_name, "file_path": model_path, "used": 0, "class_list": class_list,
                      "uuid": str(uuid.uuid4()), "model": model}
        return model_info

    def __cleanScrollLayout(self):
        children_list = self.ui.scrollWidget.children()
        for idx, children in enumerate(children_list):
            if idx != 0:
                self.ui.scrollLayout.removeWidget(children)


class ViewClass(QtCore.QThread):
    show_cur_loaded = QtCore.pyqtSignal(QtGui.QPixmap)

    def __init__(self):
        super().__init__()
        self.video_path = None
        self.frame_idx = 0
        self.parent = None

    def run(self):
        print("加载图片开始:", datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"))
        video = cv2.VideoCapture(self.video_path)
        video.set(cv2.CAP_PROP_POS_FRAMES, self.frame_idx)
        ret, frame = video.read()
        if ret:
            # 成功读取帧
            frame = self.parent.cvMatToQImage(frame)
            pixmap = QtGui.QPixmap.fromImage(frame)
            self.show_cur_loaded.emit(pixmap)
            print("加载图片完成:", datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"))
        self.exit()

    def load_img(self, parent, video_path, frame_idx):
        self.parent = parent
        self.video_path = video_path
        self.frame_idx = frame_idx
        self.run()
